from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate, MessagesPlaceholder

from langchain_demo.my_llm import llm

# ICL:
#  2 🦜 9 的结果是多少？

examples = [
    {"input": "2 🦜 2", "output": "4"},
    {"input": "2 🦜 3", "output": "6"},
]

example_template = ChatPromptTemplate.from_messages(
    [('human', '{input}'), ('ai', '{output}')],
)

few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_template,
    examples=examples,
)

prompt_template = ChatPromptTemplate.from_messages(
    [('system', '你是智能助手'),few_shot_prompt,MessagesPlaceholder('message')],
)

chain = prompt_template | llm
res = chain.invoke({"message": [HumanMessage('2 🦜 9 的结果是多少？')]})
print(res)